/******************************************************************************
	This file is a part of the Snapshot Bioinformatics Project

	Copyright (C) 2012-2013 Jarom Schow

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
******************************************************************************/

#include "SnapshotTree.h"

#include <QtCore/QTextStream>
#include <QtCore/QStack>
#include <QtCore/QDebug>

#include <QtSql/QSqlQuery>

#include "SqlInsertBuilder.h"
#include "SeqRange.h"
#include "SnapshotGeneCollection.h"

/**
**/
SnapshotNode::SnapshotNode() : m_taxId( -1 ), m_genCodeId( -1 ), m_mitoGenCodeId( -1 ), m_isVisible( true ), 
	m_parent( NULL ), m_coverage( 0.0f )
{
}

/**
**/
SnapshotNode::SnapshotNode( int taxId, const QString &rank, int genCodeId, int mitoGenCodeId ) : m_taxId( taxId ),
	m_rank( rank ), m_genCodeId( genCodeId ), m_mitoGenCodeId( mitoGenCodeId ), m_isVisible( true ), 
	m_parent( NULL ), m_coverage( 0.0f )
{
}

/**
**/
SnapshotNode::~SnapshotNode()
{
	qDeleteAll( m_children );

	m_parent = NULL;
}

/**
**/
SnapshotNode* SnapshotNode::getParent() const
{
	return m_parent;
}

/**
**/
int SnapshotNode::getNumChildren() const
{
	return m_children.count();
}

/**
**/
int SnapshotNode::indexOf( SnapshotNode *child ) const
{
	return m_children.indexOf( child );
}

/**
**/
SnapshotNode* SnapshotNode::getChild( int which ) const
{
	if( which < 0 || which >= m_children.count() )
		return NULL;

	return m_children[which];
}

/**
**/
QListIterator<SnapshotNode*> SnapshotNode::getChildIterator() const
{
	return QListIterator<SnapshotNode*>( m_children );
}

/**
**/
void SnapshotNode::addChild( SnapshotNode *child )
{
	if( !child )
		return;

	child->m_parent = this;
	m_children.append( child );
}

/**
**/
void SnapshotNode::setScientificName( const QString &name )
{
	m_scientificName = name;
}

/**
**/
QString SnapshotNode::getScientificName() const
{
	return m_scientificName;
}

/**
**/
void SnapshotNode::setCommonName( const QString &name )
{
	m_commonName = name;
}

/**
**/
QString SnapshotNode::getCommonName() const
{
	return m_commonName;
}

/**
**/
void SnapshotNode::setTaxonId( int id )
{
	m_taxId = id;
}

/**
**/
int SnapshotNode::getTaxonId() const
{
	return m_taxId;
}

/**
**/
void SnapshotNode::setGeneticCode( int id )
{
	m_genCodeId = id;
}

/**
**/
int SnapshotNode::getGeneticCode() const
{
	return m_genCodeId;
}

/**
**/
void SnapshotNode::setMitoGeneticCode( int id )
{
	m_mitoGenCodeId = id;
}

/**
**/
int SnapshotNode::getMitoGeneticCode() const
{
	return m_mitoGenCodeId;
}

/**
**/
void SnapshotNode::setRank( const QString &rank )
{
	m_rank = rank;
}

/**
**/
QString SnapshotNode::getRank() const
{
	return m_rank;
}

/**
**/
int SnapshotNode::getNumSpecies() const
{
	if( m_rank == "species" )
		return 1;

	int	numSpecies = 0;

	QListIterator<SnapshotNode*>	childIt( m_children );

	while( childIt.hasNext() )
		numSpecies += childIt.next()->getNumSpecies();

	return numSpecies;
}

/**
**/
float SnapshotNode::getCoverage() const
{
	return m_coverage;
}

/**
**/
int SnapshotNode::getNumGenes() const
{
	return m_genes.count();
}

/**
**/
SnapshotGene* SnapshotNode::getGene( int which ) const
{
	if( which < 0 || which >= m_genes.count() )
		return NULL;

	return m_genes[which];
}

/**
**/
QListIterator<SnapshotGene*> SnapshotNode::getGeneIterator() const
{
	return QListIterator<SnapshotGene*>( m_genes );
}

/**
**/
bool SnapshotNode::isVisible() const
{
	return m_isVisible;
}

/**
**/
void SnapshotNode::show( bool onOff )
{
	m_isVisible = onOff;
}

/**
**/
void SnapshotNode::updateVisibility( int visFlags )
{
	QListIterator<SnapshotNode*>	childIt( m_children );

	while( childIt.hasNext() )
		childIt.next()->updateVisibility( visFlags );

	if( (visFlags & SnapshotTree::HideUnclassified) && m_scientificName.startsWith( "unclassified" ) )
		m_isVisible = false;
	else if( (visFlags & SnapshotTree::HideEnvSamples) && m_scientificName.startsWith( "environmental samples" ) )
		m_isVisible = false;
	else if( (visFlags & SnapshotTree::HideEmptyNodes) && m_genes.isEmpty() )
		m_isVisible = false;
	else
		m_isVisible = true;
}

/**
**/
void SnapshotNode::getAllTaxa( QList<SnapshotNode*> &nodes, const QString &rank ) const
{
	if( !rank.isEmpty() ) {
		if( m_rank == rank ){
			nodes.append( const_cast<SnapshotNode*>( this ) );
			return;
		}
	}
	else {
		if( m_children.isEmpty() ){
			nodes.append( const_cast<SnapshotNode*>( this ) );
			return;
		}
	}

	QListIterator<SnapshotNode*>	childIt( m_children );

	while( childIt.hasNext() )
		childIt.next()->getAllTaxa( nodes, rank );
}

/**
**/
bool SnapshotNode::loadFromDatabase( Bio::SqlServer &server )
{
	QString		sql = QString( "SELECT name, name_class FROM taxon_name WHERE taxon_id = %1" ).arg( m_taxId );

	if( !server.exec( sql ) )
		return false;

	QSqlQuery	query = server.getQuery();

	if( query.first() ){
		do {
			QString	nameClass = query.value( 1 ).toString();

			if( nameClass == "scientific name" )
				m_scientificName = query.value( 0 ).toString();
			else if( nameClass == "genbank common name" )
				m_commonName = query.value( 0 ).toString();
			else if( nameClass == "common name" && m_commonName.isEmpty() )
				m_commonName = query.value( 0 ).toString();
		} while( query.next() );
	}

	return true;
}

/**
**/
bool SnapshotNode::loadChildrenFromDatabase( Bio::SqlServer &server, Bio::Progress *progress )
{
	QString		sql = QString( "SELECT taxon_id, node_rank, genetic_code, mito_genetic_code FROM taxon "
						"WHERE parent_taxon_id = %1" ).arg( m_taxId );
	
	qDeleteAll( m_children );
	m_children.clear();

	if( !server.exec( sql ) )
		return false;

	QSqlQuery	query = server.getQuery();

	if( !query.first() )
		return true;

	int	size = query.size();

	if( size > 0 )
		m_children.reserve( size );

	do {
		SnapshotNode	*child = new SnapshotNode( query.value( 0 ).toInt(), query.value( 1 ).toString(), 
										query.value( 2 ).toInt(), query.value( 3 ).toInt() );
		addChild( child );
		if( progress )
			progress->step();
	} while( query.next() );

	QListIterator<SnapshotNode*>	childIt( m_children );

	while( childIt.hasNext() ){
		SnapshotNode	*child = childIt.next();

		child->loadFromDatabase( server );

		if( !child->loadChildrenFromDatabase( server, progress ) )
			return false;
	}

	return true;
}

/**
**/
bool SnapshotNode::loadGenes( SnapshotGeneCollection *collection, Bio::SqlServer &server, Bio::Progress *progress )
{
	int	typeTermId = server.getTermId( "gene", "SeqFeature Keys" );
	int	termId = server.getTermId( "gene", "Annotation Tags" );

	if( termId < 0 )
		return false;

	QString		sql = QString( "SELECT bioentry_id, seqfeature_id, accession, value FROM bioentry JOIN "
						"seqfeature USING (bioentry_id) JOIN seqfeature_qualifier_value USING (seqfeature_id) "
						"WHERE taxon_id = %1 AND type_term_id = %2 AND term_id = %3" ).arg( m_taxId )
						.arg( typeTermId ).arg( termId );

	if( !server.exec( sql ) )
		return false;

	QSqlQuery			query = server.getQuery();
	QSet<SnapshotGene*>	genes;

	if( query.first() ){
		do {
			SnapshotGene	*gene = collection->getGene( query.value( 3 ).toString() );

			if( !gene )
				return false;

			Bio::SeqRange	range = server.getFeatureRange( query.value( 1 ).toInt() );

			gene->addSequence( m_scientificName, m_taxId, query.value( 0 ).toInt(), query.value( 2 ).toString(), range );
			genes.insert( gene );
		} while( query.next() );
	}

	QListIterator<SnapshotNode*>	childIt( m_children );

	while( childIt.hasNext() ){
		SnapshotNode	*child = childIt.next();

		if( !child->loadGenes( collection, server, progress ) )
			return false;

		QListIterator<SnapshotGene*>	geneIt( child->getGeneIterator() );

		while( geneIt.hasNext() )
			genes.insert( geneIt.next() );

		if( progress )
			progress->step();
	}

	m_genes = genes.toList();

	return true;
}

/**
**
void SnapshotNode::updateGeneCoverage( int numGenes )
{
	if( numGenes > 0 ){
		QListIterator<SnapshotGene*>	geneIt( m_genes );

		int	geneCount = 0;

		while( geneIt.hasNext() ){
			if( geneIt.next()->isSelected() )
				geneCount++;
		}

		m_coverage = float( geneCount ) / float( numGenes );
	}
	else
		m_coverage = 0.0f;

	QListIterator<SnapshotNode*>	childIt( m_children );

	while( childIt.hasNext() )
		childIt.next()->updateGeneCoverage( numGenes );
}

/**
**/
void SnapshotNode::updateGeneCoverage( int numGenes, int &numLeafs, float &coverage )
{
	int		localLeafs = 0;
	float	localCoverage = 0.0f;

	QListIterator<SnapshotNode*>	childIt( m_children );

	while( childIt.hasNext() )
		childIt.next()->updateGeneCoverage( numGenes, localLeafs, localCoverage );

	if( m_isVisible && numGenes > 0 ){
		if( m_rank == "species" ){
			QListIterator<SnapshotGene*>	geneIt( m_genes );

			int	geneCount = 0;

			while( geneIt.hasNext() ){
				if( geneIt.next()->isSelected() )
					geneCount++;
			}

			m_coverage = float( geneCount ) / float( numGenes );

			numLeafs += 1;
			coverage += m_coverage;
		}
		else if( localLeafs > 0 ){
			m_coverage = localCoverage / float( localLeafs );
			numLeafs += localLeafs;
			coverage += localCoverage;
		}
		else {
			m_coverage = 0.0f;
		}
	}
	else
		m_coverage = 0.0f;
}

/**
**/
SnapshotTree::SnapshotTree() : m_visFlags( ShowAll )
{
}

/**
**/
SnapshotTree::~SnapshotTree()
{
	clear();
}

/**
**/
int SnapshotTree::getNumRoots() const
{
	return m_roots.count();
}

/**
**/
int SnapshotTree::getNumSpecies() const
{
	int	numSpecies = 0;

	QListIterator<SnapshotNode*>	rootIt( m_roots );

	while( rootIt.hasNext() )
		numSpecies += rootIt.next()->getNumSpecies();

	return numSpecies;
}

/**
**/
int SnapshotTree::indexOf( SnapshotNode *root ) const
{
	return m_roots.indexOf( root );
}

/**
**/
SnapshotNode* SnapshotTree::getRoot( int which ) const
{
	if( which < 0 || which >= m_roots.count() )
		return NULL;

	return m_roots[which];
}

/**
**/
QListIterator<SnapshotNode*> SnapshotTree::getRootIterator() const
{
	return QListIterator<SnapshotNode*>( m_roots );
}

/**
**/
QList<SnapshotNode*> SnapshotTree::getAllTaxa( const QString &rank ) const
{
	QList<SnapshotNode*>			nodes;
	QListIterator<SnapshotNode*>	rootIt( m_roots );

	while( rootIt.hasNext() )
		rootIt.next()->getAllTaxa( nodes, rank );

	return nodes;
}

/**
**/
bool SnapshotTree::addFromDatabase( int rootTaxonId, Bio::SqlServer &server, SnapshotGeneCollection *collection )
{
	QString		sql = QString( "SELECT node_rank, genetic_code, mito_genetic_code, left_value, right_value FROM taxon "
						"WHERE taxon_id = %1" ).arg( rootTaxonId );

	if( !server.exec( sql ) )
		return false;

	QSqlQuery	query = server.getQuery();

	if( !query.first() )
		return false;

	SnapshotNode	*root = new SnapshotNode( rootTaxonId, query.value( 0 ).toString(), 
									query.value( 1 ).toInt(), query.value( 2 ).toInt() );

	Bio::Progress	progress( "Loading taxa from database...", query.value( 4 ).toInt() - query.value( 3 ).toInt() );

	root->loadFromDatabase( server );
	
	if( !root->loadChildrenFromDatabase( server, &progress ) ){
		delete root;
		return false;
	}

	if( !root->loadGenes( collection, server, &progress ) ){
		delete root;
		return false;
	}

	m_roots.append( root );

	updateGeneCoverage( collection->getNumGenes() );
	updateVisibility();

	emit rootAdded( root );

	progress.finish();

	return true;
}

/**
**/
void SnapshotTree::updateGeneCoverage( int numGenes )
{
	int		numLeafs = 0;
	float	coverage = 0.0f;

	QListIterator<SnapshotNode*>	rootIt( m_roots );

	while( rootIt.hasNext() )
		rootIt.next()->updateGeneCoverage( numGenes, numLeafs, coverage );

	emit coverageChanged();
}

/**
**/
bool SnapshotTree::getIgnoreEmpty() const
{
	return (m_visFlags & HideEmptyNodes) != 0;
}

/**
**/
void SnapshotTree::setIgnoreEmpty( bool onOff )
{
	if( onOff == getIgnoreEmpty() )
		return;

	if( onOff )
		m_visFlags |= HideEmptyNodes;
	else
		m_visFlags &= ~HideEmptyNodes;

	updateVisibility();
	emit visibilityChanged();
}

/**
**/
bool SnapshotTree::getIgnoreUnclassified() const
{
	return (m_visFlags & HideUnclassified) != 0;
}

/**
**/
void SnapshotTree::setIgnoreUnclassified( bool onOff )
{
	if( onOff == getIgnoreUnclassified() )
		return;

	if( onOff )
		m_visFlags |= HideUnclassified;
	else
		m_visFlags &= ~HideUnclassified;

	updateVisibility();
	emit visibilityChanged();
}

/**
**/
bool SnapshotTree::getIgnoreEnvSamp() const
{
	return (m_visFlags & HideEnvSamples) != 0;
}

/**
**/
void SnapshotTree::setIgnoreEnvSamp( bool onOff )
{
	if( onOff == getIgnoreEnvSamp() )
		return;

	if( onOff )
		m_visFlags |= HideEnvSamples;
	else
		m_visFlags &= ~HideEnvSamples;

	updateVisibility();
	emit visibilityChanged();
}

/**
**/
void SnapshotTree::updateVisibility()
{
	QListIterator<SnapshotNode*>	rootIt( m_roots );

	while( rootIt.hasNext() )
		rootIt.next()->updateVisibility( m_visFlags );
}

/**
**/
void SnapshotTree::clear()
{
	qDeleteAll( m_roots );
	m_roots.clear();
	emit cleared();
}
